import torch
import numpy as np

N, D_in, H, D_out = 64, 1000, 100, 10  # 输入，维度，中间层，输出
# 随机创建一些训练数据
x = torch.randn(N, D_in)
y = torch.randn(N, D_out)

w1 = torch.randn(D_in, H, requires_grad=True)
w2 = torch.randn(H, D_out, requires_grad=True)

leaning_rate = 1e-6
for it in range(500):
    # Forward pass
    y_pred = x.mm(w1).clamp(min=0).mm(w2)  # N*D_out

    # compute loss
    loss = (y_pred-y).pow(2).sum()  # MSE loss 均方误差
    print(it, loss.item())

    # Backward pass
    loss.backward()

    # Update weights of w1 and w2
    with torch.no_grad():
        w1 -= leaning_rate * w1.grad
        w2 -= leaning_rate * w2.grad
        w1.grad.zero_()
        w2.grad.zero_()